package mr

import "log"
import "net"
import "os"
import "net/rpc"
import "net/http"
import "sync"
import "strconv"

type pair struct {
    a int
    b int
}


type Coordinator struct {
    nReduce int
    nMap int

    files []string
    map_state []pair
    reduce_state []pair
    map_complete bool

    sv_done bool

    files_lock sync.Mutex
    map_state_lock sync.Mutex
    reduce_state_lock sync.Mutex
}

// Your code here -- RPC handlers for the worker to call.

func (c *Coordinator) init(files *[]string, nReduce int) {
    c.files = (*files)
    c.nReduce = nReduce
    nMap := len(*files)
    c.nMap = nMap
    for i := 1; i <= nMap; i++ {
        c.map_state = append(c.map_state, pair{0, -1})
    }
    for i := 1; i <= nReduce; i++ {
        c.reduce_state = append(c.reduce_state, pair{0, -1})
    }
    // fmt.Println("Master initialized.")
}

func (c *Coordinator) assign_map_job(pid int) int {
    if c.map_complete {
        return -1;
    }
    c.map_state_lock.Lock()
    for i := 0; i < c.nMap; i++ {
        if c.map_state[i].a == 0 {
            c.map_state[i].a = 1
            c.map_state[i].b = pid
            c.map_state_lock.Unlock()
            return i
        } else if c.map_state[i].a == 1 {
            args := true
            reply := false
            if !c.call("Ping.Tic", c.map_state[i].b, &args, &reply) {
                //worker not working
                c.map_state[i].b = pid
                c.map_state_lock.Unlock()
                return i
            }
        }
    }
    c.map_state_lock.Unlock()
    return -1
}

func (c *Coordinator) assign_reduce_job(pid int) int {
    if !c.map_complete {
        c.map_state_lock.Lock()
        mark := true
        for i := 0; i < c.nMap; i++ {
            if c.map_state[i].a != 2 {
                mark = false
                break;
            }
        }
        c.map_state_lock.Unlock()
        c.map_complete = mark
    }
    if !c.map_complete {
        return -1
    }
    c.reduce_state_lock.Lock()
    for i := 0; i < c.nReduce; i++ {
        if c.reduce_state[i].a == 0 {
            c.reduce_state[i].a = 1
            c.reduce_state[i].b = pid
            c.reduce_state_lock.Unlock()
            return i
        } else if c.reduce_state[i].a == 1 {
            args := true
            reply := false
            if !c.call("Ping.Tic", c.reduce_state[i].b, &args, &reply) {
                c.reduce_state[i].b = pid
                c.reduce_state_lock.Unlock()
                return i
            }
        }
    }
    c.reduce_state_lock.Unlock()
    return -1
}

func (c *Coordinator) get_map_file(id int) string {
    ret := string("")
    c.files_lock.Lock()
    ret = c.files[id]
    c.files_lock.Unlock()
    return ret
}

func (c *Coordinator) assign_new_job(pid int) (int, string, int) {

    if c.Done() {
        return -1, "ph", -1
    }
    t := c.assign_map_job(pid)
    if t != -1 {
        return t, c.get_map_file(t), 0
    } else {
        t = c.assign_reduce_job(pid)
        if t != -1 {
            return t, strconv.Itoa(t), 1
        }
    }
    return -1, "ph", -1
}

func (c* Coordinator) commit_map(job_id int) bool {
    c.map_state_lock.Lock()
    if c.map_state[job_id].a == 1 {
        c.map_state[job_id].a = 2
        c.map_state_lock.Unlock()
        return true
    }
    c.map_state_lock.Unlock()
    return false
}

func (c *Coordinator) commit_reduce(job_id int) bool {
    c.reduce_state_lock.Lock()
    if c.reduce_state[job_id].a == 1 {
        c.reduce_state[job_id].a = 2
        c.reduce_state_lock.Unlock()
        return true
    }
    c.reduce_state_lock.Unlock()
    return false
}


// RPC methods

func (c* Coordinator) AssignJob(args *AssignJobArgs, reply *AssignJobReply) error {
    Id, Name, Job := c.assign_new_job(args.Pid)
    reply.Id = Id
    reply.Name = Name
    reply.Job = Job
    return nil
}

func (c* Coordinator) GetBasics(args *int, reply *PairReply) error {
    reply.A = c.nMap
    reply.B = c.nReduce
    return nil
}

func (c* Coordinator) MapCommit(args *int, reply *bool) error {
    (*reply) = c.commit_map(*args)
    return nil
}

func (c* Coordinator) ReduceCommit(args *int, reply *bool) error {
    (*reply) = c.commit_reduce(*args)
    return nil
}

//
// an example RPC handler.
//
// the RPC argument and reply types are defined in rpc.go.
//
func (c *Coordinator) Example(args *ExampleArgs, reply *ExampleReply) error {
    reply.Y = args.X + 1
    return nil
}

//
// start a thread that listens for RPCs from worker.go
//
func (c *Coordinator) server() {
    rpc.Register(c)
    rpc.HandleHTTP()
    //l, e := net.Listen("tcp", ":1234")
    sockname := coordinatorSock()
    os.Remove(sockname)
    l, e := net.Listen("unix", sockname)
    if e != nil {
        log.Fatal("listen error:", e)
    }
    go http.Serve(l, nil)
}

//
// main/mrcoordinator.go calls Done() periodically to find out
// if the entire job has finished.
//
func (c *Coordinator) Done() bool {
    ret := true
    if c.sv_done {
        return ret
    }
    for i := 0; i < c.nReduce; i++ {
        if c.reduce_state[i].a != 2 {
            ret = false
        }
    }
    c.sv_done = ret
    return ret
}

//
// create a Coordinator.
// main/mrcoordinator.go calls this function.
// nReduce is the number of reduce tasks to use.
//
func MakeCoordinator(files []string, nReduce int) *Coordinator {
    c := Coordinator{}

    // Your code here.
    c.init(&files, nReduce)

    c.server()
    return &c
}

func (c *Coordinator) call(rpcname string, pid int, args interface{}, reply interface{}) bool {
    // c, err := rpc.DialHTTP("tcp", "127.0.0.1"+":1234")
    sockname := cookupSock(pid)
    cc, err := rpc.DialHTTP("unix", sockname)
    if err != nil {
        log.Println("dialing:", err)
        return false
    }
    defer cc.Close()

    err = cc.Call(rpcname, args, reply)
    if err == nil {
        return true
    }

    log.Println(err)
    return false
}